import os
from PIL import Image
from torchvision import transforms
from tqdm import tqdm  # 用于显示进度条

# 设置原始数据集路径和增强数据集的保存路径
original_data_dir = r'E:\Desktop\实验六\Monkeys'  # 原始数据集路径
augmented_data_dir = r'E:\Desktop\实验六\DA_Monkeys'  # 增强后的数据集保存路径

# 定义数据增强操作，每个操作生成一张增强图片
augmentations = [
    transforms.RandomHorizontalFlip(p=1.0),  # 水平翻转
    transforms.RandomRotation(10),  # 旋转±10度
    transforms.ColorJitter(brightness=0.5),  # 亮度调整
]

# 确保增强数据集的保存路径存在
os.makedirs(augmented_data_dir, exist_ok=True)

# 遍历每个类别的文件夹
for category in os.listdir(original_data_dir):
    category_path = os.path.join(original_data_dir, category)
    augmented_category_path = os.path.join(augmented_data_dir, category)
    os.makedirs(augmented_category_path, exist_ok=True)

    # 遍历该类别下的每张图像
    for image_name in tqdm(os.listdir(category_path), desc=f"Processing {category}"):
        image_path = os.path.join(category_path, image_name)

        # 打开图像
        with Image.open(image_path).convert("RGB") as img:
            # 保存原图
            img.save(os.path.join(augmented_category_path, image_name))

            # 对每个增强操作生成并保存增强后的图像
            for i, augmentation in enumerate(augmentations):
                augmented_img = augmentation(img)  # 应用数据增强操作
                augmented_img_name = f"{os.path.splitext(image_name)[0]}_aug_{i}.jpg"
                augmented_img_path = os.path.join(augmented_category_path, augmented_img_name)
                augmented_img.save(augmented_img_path)
